#ifndef EXP__GRAPH_HPP
#define EXP__GRAPH_HPP

#include <iostream>
#include <vector>
#include <unordered_set>
#include <queue>
#include <stack>

namespace my_data
{
/**
 * @brief 无向连通图类
 */ 
class Graph
{
public:
/**
 * @brief 构造函数
 */ 
  explicit Graph(const int &num_vertices);
/**
 * @brief 添加一条从连接结点u,v的边
 */ 
  void addEdge(const int &u, const int &v);
/**
 * @brief 深度优先遍历
 * @return 生成树集合
 */ 
  std::unordered_set<int> DFS(const int &start_vertex);
/**
 * @brief 广度优先遍历
 * @return 生成树集合
 */ 
  std::unordered_set<int> BFS(const int &start_vertex);
/**
 * @brief 设置顶点的名字
 */ 
  void setNames(const std::vector<std::string> &names);

private:
  int num_vertices_; // 顶点数
  std::vector<std::vector<int>> adj_list_; // 邻接表
  std::vector<std::string> names_;
};


Graph::Graph(const int &num_vertices)
:num_vertices_(num_vertices)
{
  adj_list_.resize(num_vertices_);
}

void Graph::setNames(const std::vector<std::string> &names)
{
  names_ = names;
}

void Graph::addEdge(const int &u, const int &v)
{
  adj_list_[u].emplace_back(v);
  adj_list_[v].emplace_back(u);
}

std::unordered_set<int> Graph::DFS(const int &start_vertex)
{
  std::unordered_set<int> visited;
  std::unordered_set<int> tree;
  std::stack<int> stack;
  // 将起点加入visited,tree和stack中
  visited.insert(start_vertex);
  tree.insert(start_vertex);
  stack.push(start_vertex);
  while(!stack.empty())
  {
    //取出栈顶
    int vertex = stack.top();
    if(!names_.empty())
    {
      std::cout<<vertex<<":"<<names_[vertex-1]<<" ";
    }
    else
    {
      std::cout<<vertex<<" ";
    }
    stack.pop();
    // 遍历当前结点的所有邻接结点
    for(int neighbor : adj_list_[vertex])
    {
      // 如果邻接结点没有被访问，将邻接结点加入visited和tree中，并将邻接结点入栈
      if (visited.find(neighbor) == visited.end()) 
      {
        visited.insert(neighbor);
        tree.insert(neighbor);
        stack.push(neighbor);
      } 
    }
  }
  return tree;
}

std::unordered_set<int> Graph::BFS(const int &start_vertex)
{
  std::unordered_set<int> visited;
  std::unordered_set<int> tree;
  std::queue<int> queue;
  // 将起点加入visited,tree和queue中
  visited.insert(start_vertex);
  tree.insert(start_vertex);
  queue.push(start_vertex);
  // 当队列不为空时，继续遍历
  while(!queue.empty())
  {
    // 取出队列头
    int vertex = queue.front();
    if(!names_.empty())
    {
      std::cout<<vertex<<":"<<names_[vertex-1]<<" ";
    }
    else
    {
      std::cout<<vertex<<" ";
    }
    queue.pop();
    // 遍历当前结点的所有邻接结点
    for(int neighbor : adj_list_[vertex])
    {
      // 如果邻接结点没有被访问，将邻接结点加入visited和tree中，并将邻接结点加入队列
      if (visited.find(neighbor) == visited.end()) 
      {
        visited.insert(neighbor);
        tree.insert(neighbor);
        queue.push(neighbor);
      } 
    }
  }
  return tree;
}


} //namespace my_data

#endif //EXP__GRAPH_HPP



